# Copyright 2024 The KServe Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pytest import approx


bert_token_classification_return_raw_logits_expected_output = {
    "predictions": [
        [
            {
                0: approx(10.125612258911133, abs=0.000009),
                1: approx(-1.818464756011963, abs=0.000009),
                2: approx(-1.3191171884536743, abs=0.000009),
                3: approx(-1.9324339628219604, abs=0.000009),
                4: approx(-1.4850239753723145, abs=0.000009),
                5: approx(-1.69266676902771, abs=0.000009),
                6: approx(-0.898107647895813, abs=0.000009),
                7: approx(-1.725127935409546, abs=0.000009),
                8: approx(0.32057392597198486, abs=0.000009),
            },
            {
                0: approx(9.647306442260742, abs=0.000009),
                1: approx(-2.112745523452759, abs=0.000009),
                2: approx(-0.8831668496131897, abs=0.000009),
                3: approx(-2.719135046005249, abs=0.000009),
                4: approx(-0.47473397850990295, abs=0.000009),
                5: approx(-2.2424018383026123, abs=0.000009),
                6: approx(0.6101505160331726, abs=0.000009),
                7: approx(-2.2192084789276123, abs=0.000009),
                8: approx(-0.654518187046051, abs=0.000009),
            },
            {
                0: approx(10.364561080932617, abs=0.000009),
                1: approx(-2.240158796310425, abs=0.000009),
                2: approx(-0.9236820340156555, abs=0.000009),
                3: approx(-2.623302936553955, abs=0.000009),
                4: approx(-0.501063346862793, abs=0.000009),
                5: approx(-1.9418426752090454, abs=0.000009),
                6: approx(-0.041013482958078384, abs=0.000009),
                7: approx(-2.12089204788208, abs=0.000009),
                8: approx(-1.2565152645111084, abs=0.000009),
            },
            {
                0: approx(9.21961784362793, abs=0.000009),
                1: approx(-2.1359012126922607, abs=0.000009),
                2: approx(-0.1689995676279068, abs=0.000009),
                3: approx(-3.0277929306030273, abs=0.000009),
                4: approx(0.2589835822582245, abs=0.000009),
                5: approx(-2.4426753520965576, abs=0.000009),
                6: approx(0.4815778136253357, abs=0.000009),
                7: approx(-2.3223936557769775, abs=0.000009),
                8: approx(-0.23837946355342865, abs=0.000009),
            },
            {
                0: approx(11.023712158203125, abs=0.000009),
                1: approx(-2.0757784843444824, abs=0.000009),
                2: approx(-0.8648976683616638, abs=0.000009),
                3: approx(-2.6088907718658447, abs=0.000009),
                4: approx(-0.781069278717041, abs=0.000009),
                5: approx(-1.7040152549743652, abs=0.000009),
                6: approx(-0.5265535712242126, abs=0.000009),
                7: approx(-1.7141871452331543, abs=0.000009),
                8: approx(-1.0127886533737183, abs=0.000009),
            },
            {
                0: approx(9.288057327270508, abs=0.000009),
                1: approx(-1.9795234203338623, abs=0.000009),
                2: approx(0.3095782697200775, abs=0.000009),
                3: approx(-2.781409740447998, abs=0.000009),
                4: approx(-0.4492073655128479, abs=0.000009),
                5: approx(-1.7794851064682007, abs=0.000009),
                6: approx(-0.3115033507347107, abs=0.000009),
                7: approx(-2.0755014419555664, abs=0.000009),
                8: approx(-0.8829556107521057, abs=0.000009),
            },
            {
                0: approx(7.968824863433838, abs=0.000009),
                1: approx(-2.0212578773498535, abs=0.000009),
                2: approx(-0.43740716576576233, abs=0.000009),
                3: approx(-3.6762948036193848, abs=0.000009),
                4: approx(-0.012759938836097717, abs=0.000009),
                5: approx(-2.320978879928589, abs=0.000009),
                6: approx(0.6871926784515381, abs=0.000009),
                7: approx(-2.541991949081421, abs=0.000009),
                8: approx(0.32606029510498047, abs=0.000009),
            },
        ],
        [
            {
                0: approx(10.125612258911133, abs=0.000009),
                1: approx(-1.818464756011963, abs=0.000009),
                2: approx(-1.3191171884536743, abs=0.000009),
                3: approx(-1.9324339628219604, abs=0.000009),
                4: approx(-1.4850239753723145, abs=0.000009),
                5: approx(-1.69266676902771, abs=0.000009),
                6: approx(-0.898107647895813, abs=0.000009),
                7: approx(-1.725127935409546, abs=0.000009),
                8: approx(0.32057392597198486, abs=0.000009),
            },
            {
                0: approx(9.647306442260742, abs=0.000009),
                1: approx(-2.112745523452759, abs=0.000009),
                2: approx(-0.8831668496131897, abs=0.000009),
                3: approx(-2.719135046005249, abs=0.000009),
                4: approx(-0.47473397850990295, abs=0.000009),
                5: approx(-2.2424018383026123, abs=0.000009),
                6: approx(0.6101505160331726, abs=0.000009),
                7: approx(-2.2192084789276123, abs=0.000009),
                8: approx(-0.654518187046051, abs=0.000009),
            },
            {
                0: approx(10.364561080932617, abs=0.000009),
                1: approx(-2.240158796310425, abs=0.000009),
                2: approx(-0.9236820340156555, abs=0.000009),
                3: approx(-2.623302936553955, abs=0.000009),
                4: approx(-0.501063346862793, abs=0.000009),
                5: approx(-1.9418426752090454, abs=0.000009),
                6: approx(-0.041013482958078384, abs=0.000009),
                7: approx(-2.12089204788208, abs=0.000009),
                8: approx(-1.2565152645111084, abs=0.000009),
            },
            {
                0: approx(9.21961784362793, abs=0.000009),
                1: approx(-2.1359012126922607, abs=0.000009),
                2: approx(-0.1689995676279068, abs=0.000009),
                3: approx(-3.0277929306030273, abs=0.000009),
                4: approx(0.2589835822582245, abs=0.000009),
                5: approx(-2.4426753520965576, abs=0.000009),
                6: approx(0.4815778136253357, abs=0.000009),
                7: approx(-2.3223936557769775, abs=0.000009),
                8: approx(-0.23837946355342865, abs=0.000009),
            },
            {
                0: approx(11.023712158203125, abs=0.000009),
                1: approx(-2.0757784843444824, abs=0.000009),
                2: approx(-0.8648976683616638, abs=0.000009),
                3: approx(-2.6088907718658447, abs=0.000009),
                4: approx(-0.781069278717041, abs=0.000009),
                5: approx(-1.7040152549743652, abs=0.000009),
                6: approx(-0.5265535712242126, abs=0.000009),
                7: approx(-1.7141871452331543, abs=0.000009),
                8: approx(-1.0127886533737183, abs=0.000009),
            },
            {
                0: approx(9.288057327270508, abs=0.000009),
                1: approx(-1.9795234203338623, abs=0.000009),
                2: approx(0.3095782697200775, abs=0.000009),
                3: approx(-2.781409740447998, abs=0.000009),
                4: approx(-0.4492073655128479, abs=0.000009),
                5: approx(-1.7794851064682007, abs=0.000009),
                6: approx(-0.3115033507347107, abs=0.000009),
                7: approx(-2.0755014419555664, abs=0.000009),
                8: approx(-0.8829556107521057, abs=0.000009),
            },
            {
                0: approx(7.968824863433838, abs=0.000009),
                1: approx(-2.0212578773498535, abs=0.000009),
                2: approx(-0.43740716576576233, abs=0.000009),
                3: approx(-3.6762948036193848, abs=0.000009),
                4: approx(-0.012759938836097717, abs=0.000009),
                5: approx(-2.320978879928589, abs=0.000009),
                6: approx(0.6871926784515381, abs=0.000009),
                7: approx(-2.541991949081421, abs=0.000009),
                8: approx(0.32606029510498047, abs=0.000009),
            },
        ],
    ]
}
